package com.wdemo.service.impl;

import cn.hutool.core.util.ObjectUtil;
import com.wdemo.config.AIProperties;
import com.wdemo.enums.ChatTypeEnum;
import com.wdemo.memory.ChatMemoryManager;
import com.wdemo.service.ChatService;
import com.wdemo.service.ChatSessionService;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;

import static org.springframework.ai.chat.memory.ChatMemory.CONVERSATION_ID;

/**
 * @ClassName chatServiceImpl
 * @Description TODO
 * @Author WDEMO
 * @Date 2025/8/29 12:25
 */
@Service
@RequiredArgsConstructor
public class ChatServiceImpl implements ChatService {

    private final ChatClient chatClient;


    private final AIProperties aiProperties;

    private final VectorStore redisVectorStore;

    private final ChatSessionService chatSessionService;

    private final ChatMemoryManager chatMemoryManager;

//    private final ServeItemTool serveItemTool;

    private final ToolCallbackProvider toolCallbackProvider;

    // 存储大模型的生成状态，这里采用ConcurrentHashMap是确保线程安全
    // 目前的版本暂时用Map实现，如果考虑分布式环境的话，可以考虑用redis来实现
    private static final Map<String, Boolean> GENERATE_STATUS = new ConcurrentHashMap<String, Boolean>();


    @Override
    public Flux<String> chatStream(String prompt, String chatId) {

        this.updateTitle(chatId, prompt);

        return chatClient.prompt()
                .user(prompt)
                .advisors(advisorSpec -> advisorSpec.param(CONVERSATION_ID, chatId))
                .stream()
                .content()
                .concatWith(Flux.just("[END]"));

    }

    @Override
    public Flux<String> game(String prompt, String sessionId) {
        return chatClient.prompt()
                .system(aiProperties.getSystemGameMessage().get())
                .user(prompt)
                .advisors(advisorSpec -> advisorSpec.param(CONVERSATION_ID, sessionId))
                .stream()
                .content()
                // 在流结束时添加结束标记
                .concatWith(Flux.just("[END]"));
    }


    @Override
    public Flux<String> service(String prompt, String chatId) {

        //当promp不为空，且不是“你好”时，更新会话标题
        if (ObjectUtil.isNotNull(prompt) && !(prompt.trim().equals("你好"))){
            this.updateTitle(chatId, prompt);
        }

        SearchRequest searchRequest = SearchRequest.builder().query(prompt).topK(5).build();

        // 根据类型获取对应的ChatMemory
        ChatMemory chatMemory = getChatMemoryByType(ChatTypeEnum.SERVICE.getType());

        return chatClient.prompt()
                .system(aiProperties.getSystemServiceMessage().get())
                .user(prompt)
                .advisors(advisorSpec -> advisorSpec.param(CONVERSATION_ID, chatId)
                        .advisors(QuestionAnswerAdvisor.builder(redisVectorStore)
                                .searchRequest(searchRequest)
                                .build(),
                                MessageChatMemoryAdvisor.builder(chatMemory).build()))
                .toolCallbacks(toolCallbackProvider)
                .stream()
                .chatResponse()
                .doFirst(() -> GENERATE_STATUS.put(chatId, true))
                .doOnComplete(() -> GENERATE_STATUS.remove(chatId, false))
                .doOnError(throwable -> GENERATE_STATUS.remove(chatId, false))
                // 输出过程中，判断是否正在输出，如果正在输出，则继续输出，否则结束输出
                .takeWhile(chatResponse -> Optional.ofNullable(GENERATE_STATUS.get(chatId)).orElse(false))
                .map(chatResponse -> {
                    return chatResponse.getResult().getOutput().getText();
                })
                // 在流结束时添加结束标记
                .concatWith(Flux.just("[END]"));

    }

    /**
     * 更新会话标题
     * @param session 会话id
     * @param title 会话标题
     */
    private void updateTitle(String session, String title){
        chatSessionService.update(session,title,1L);
    }


    /**
     * 根据聊天类型获取对应的ChatMemory
     * @param type 聊天类型 (service, chat, pdf)
     * @return 对应的ChatMemory实例
     */
    public ChatMemory getChatMemoryByType(String type) {
        ChatTypeEnum chatType = ChatTypeEnum.fromString(type);
        return chatMemoryManager.getChatMemory(chatType);
    }
}
